from __future__ import annotations

import json
import datetime
from typing import Optional, Tuple, Type, Union, List, OrderedDict
import enum

from peewee import IntegerField
import peewee
from peewee import fn

from common.typing_utils import numeric
from common.datetime_utils import datetimestr_to_datetime, utcnow, FMT_DATETIME
from common.geometry.geom import Geom
from common.geoimage.virtual_raster import VRTScene
from common.logger import logger
from catalog.models.base_stac_item import BaseSTACItem
from catalog.utils import interpret_protocol


class ItemStatus(int, enum.Enum):
    """Indicates an item's status during the life-cycle of processing."""

    # NOTE Completed but verson low
    UNSUBMITTED = 0
    SUBMITTED = 1
    TRANSFERRING = 2
    TRANSFERRED = 3
    PROCESSING = 4
    COMPLETED = 5
    FAILED = 6
    ARCHIVED = 7


class ItemVersion(int, enum.Enum):
    """Version control of on-demand items

    Parameters
    ----------
    int : integer
        Version index is the integer type increasing by 1
    enum : str
        Version descroptions

    Notes
    -----
    NOTE every new data collection has their customized ItemVersion.
    """

    # Initial value of item indexing
    UNPROCESSED = 0

    # Production version in use
    PRODUCTION = 1


class BaseSTACItemOnDemand(BaseSTACItem):
    """Base class for all ON-DEMAND STAC Items managed by the catalog."""

    collection_id = "OnDemandDataCollection"
    production_version = 0  # Item Production Version for the whole collection
    status = IntegerField(
        index=True,
        help_text="The Status of the item. Check enums of ItemStatus",
    )
    version = IntegerField(
        index=True,
        help_text="The Version of Processing Algorithm. Check enums of ItemVersion",
    )

    @classmethod
    def construct_query(
        cls: Type[BaseSTACItem],
        spatial: Optional[Geom] = None,
        temporal: Optional[
            Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]
        ] = None,
        include_invalidated: bool = False,
        status: Optional[Union[ItemStatus, int, List[ItemStatus], List[int]]] = None,
        version: Optional[Union[ItemVersion, int, List[ItemVersion], List[int]]] = None,
    ) -> Tuple[peewee.Query, int]:
        """Construct a query for the given spatial and temporal constraints.

        Parameters
        ----------
        spatial : Optional[Geom]
            Area of interest represented by a Geom object, by default None not to filter by spatial area.
        temporal : Optional[Tuple[Optional[datetime.datetime], Optional[datetime.datetime]] ], optional
            Temporal range of interest, as a tuple of start and end datetime, or datetime string,
            by default None not to filter by temporal range.
        include_invalidated : bool
            Whether to include invalidated items in the query result, by default False not to include.
        status: Optional[Union[ItemStatus, int, List[ItemStatus], List[int]]]
            The mark indicating an item's status during the life-cycle of processing,
            by default None
        version: Optional[Union[ItemVersion, int, List[ItemVersion], List[int]]]
            The algorithm version used for processing items.
        Returns
        -------
        Tuple[peewee.Query, int]
            A query object that can be further tuned or executed, and the number of criterion that were applied.
        """

        query = cls.select()
        if not include_invalidated:
            query = query.where(cls.valid_until.is_null())

        n_conditions = 0

        # spatial queries
        if spatial is not None:
            query = query.where(
                fn.ST_Intersects(
                    cls.geom, fn.ST_GeomFromGeoJSON(spatial.to_geometry_json())
                )
            )
            n_conditions += 1

        # temporal queries for captured time
        if temporal is not None:
            start_datetime, end_datetime = temporal
            if start_datetime is not None:
                query = query.where(cls.start_datetime >= start_datetime)
                n_conditions += 1
            if end_datetime is not None:
                query = query.where(cls.end_datetime < end_datetime)
                n_conditions += 1

        # item processing status
        if status is not None:
            if isinstance(status, (ItemStatus, int)):
                query = query.where(cls.status == status)
            elif isinstance(status, list):
                query = query.where(cls.status.in_(status))
            n_conditions += 1

        # item production version
        if version is not None:
            if isinstance(version, (ItemVersion, int)):
                query = query.where(cls.version == version)
            elif isinstance(version, list):
                query = query.where(cls.version.in_(version))
            n_conditions += 1

        return query, n_conditions

    def choose_assets(
        self,
        lst_assets: List[str],
        dest_nodata: Union[Optional[numeric], List[Optional[numeric]]] = 0,
    ) -> VRTScene:
        """Create a virtual raster scene from the STAC item, by choosing assets from the given list of assets.

        Parameters
        ----------
        lst_assets : List[str]
            A list of assets to be used.
        dest_nodata : Union[Optional[numeric], List[Optional[numeric]]], optional
            Nodata value(s) to be used, by default 0

        Returns
        -------
        VRTScene
            The created virtual raster scene.
        """
        status = next(status for status in ItemStatus if status == self.status)
        if status == ItemStatus.COMPLETED:
            if self.version != self.production_version:
                logger.warning(
                    f"Item <{self.uid}> not in PRODUCTION version, you may request reprocessing."
                )
            selected_assets = OrderedDict()
            for asset_name in lst_assets:
                selected_assets[asset_name] = interpret_protocol(
                    self.info["assets"][asset_name]["href"]
                )
            return VRTScene.from_assets(selected_assets, dest_nodata=dest_nodata)
        elif status in [
            ItemStatus.SUBMITTED,
            ItemStatus.TRANSFERRING,
            ItemStatus.TRANSFERRED,
            ItemStatus.PROCESSING,
        ]:
            logger.warning(
                f"Item <{self.uid}> in processing with status: {ItemStatus.__str__(status)}. "
            )
            return None
        elif status in [ItemStatus.UNSUBMITTED, ItemStatus.FAILED]:
            logger.warning(
                f"Item <{self.uid}> not processed yet with status "
                f"{ItemStatus.__str__(status)}. Please request processing."
            )
            return None
        else:
            logger.error(f"Unknown status for item <{self.uid}>: {status}.")
            raise RuntimeError

    @classmethod
    def update_item_info(cls, item, timestamp=None, dry_run=True):
        """Update item meta info in the catalog.

        Parameters
        ----------
        item: BaseSTACItemOnDemand
            A BaseSTACItemOnDemand item.
        timestamp: datetime.datetime
            A datetime.datetime recording the updated timestamp.
        dry_run: bool
            If in dry run mode. Default: True
        """
        if item is None:
            logger.warning("Input `item` is None, skipped.")
            return

        if dry_run:
            logger.info(f"Dry run. Not update item status for <{item.uid}>")
            return

        if timestamp is None:
            timestamp = utcnow()
        timestamp_str = timestamp.strftime(FMT_DATETIME)

        item.updated = timestamp
        item.info["properties"]["updated"] = timestamp_str

        cls.update_one_item(item)
        return item

    @classmethod
    def update_item_status(
        cls, item, status, incremental=True, timestamp=None, dry_run=True
    ):
        """Update item status in the catalog.

        Parameters
        ----------
        item: BaseSTACItemOnDemand
            A BaseSTACItemOnDemand item.
        status: ItemStatus or int
            Processing status of an indexed item.
        incremental: bool
            If only allow updating item status to a higher one. Default: True
        timestamp: datetime.datetime
            A datetime.datetime recording the updated timestamp.
        dry_run: bool
            If in dry run mode. Default: True
        """
        if item is None:
            logger.warning("Input `item` is None, skipped.")
            return

        if dry_run:
            logger.info(f"Dry run. Not update item status for <{item.uid}>")
            return

        if incremental and item.status != ItemStatus.FAILED and status < item.status:
            logger.warning(
                f"Item status not updated. Must update to a higher status. "
                f"current: {item.status}, target: {status}."
            )
            return

        item.status = status
        item.info["properties"]["paii-misc:status"] = status

        if timestamp is None:
            timestamp = utcnow()
        timestamp_str = timestamp.strftime(FMT_DATETIME)

        item.updated = timestamp
        item.info["properties"]["updated"] = timestamp_str

        if status == ItemStatus.TRANSFERRED:
            item.transferred = timestamp
            item.info["properties"]["paii-lc:transferred"] = timestamp_str
        elif status == ItemStatus.COMPLETED:
            item.processed = timestamp
            item.info["properties"]["paii-lc:processed"] = timestamp_str
        cls.update_one_item(item)
        return item

    @classmethod
    def update_item_version(
        cls, item, version, incremental=True, timestamp=None, dry_run=True
    ):
        """Update item status in the catalog.

        Parameters
        ----------
        item: BaseSTACItemOnDemand
            A BaseSTACItemOnDemand item.
        version: ItemVersion or int
            Product version of an indexed item.
        incremental: bool
            If only allow updating item version to a higher one. Default: True
        timestamp: datetime.datetime
            A datetime.datetime recording the updated timestamp.
        dry_run: bool
            If in dry run mode. Default: True
        """
        if item is None:
            logger.warning("Input `item` is None, skipped.")
            return

        if dry_run:
            logger.info(f"Dry run. Not update item version for <{item.uid}>")
            return

        if incremental and version < item.version:
            logger.warning(
                f"Item version not updated. Must update to a higher version. "
                f"current: {item.version}, target: {version}."
            )
            return

        if timestamp is None:
            timestamp = utcnow()
        timestamp_str = timestamp.strftime(FMT_DATETIME)

        item.updated = timestamp
        item.info["properties"]["updated"] = timestamp_str

        item.version = version
        item.info["properties"]["paii-misc:version"] = version
        cls.update_one_item(item)
        return item


def extract_base_info(info_item):
    """Extract the base info required for an ON-DEMAND STAC Item from the given information dictionary."""
    uid = info_item["id"]
    acquired_start = datetimestr_to_datetime(
        info_item["properties"]["datetime"]["start_datetime"]
    )
    acquired_end = datetimestr_to_datetime(
        info_item["properties"]["datetime"]["end_datetime"]
    )

    processed = datetimestr_to_datetime(info_item["properties"]["paii-lc:processed"])
    published = datetimestr_to_datetime(info_item["properties"]["published"])
    indexed = datetimestr_to_datetime(info_item["properties"]["paii-lc:indexed"])
    transferred = datetimestr_to_datetime(
        info_item["properties"]["paii-lc:transferred"]
    )
    archived = datetimestr_to_datetime(info_item["properties"]["paii-lc:archived"])
    created = datetimestr_to_datetime(info_item["properties"]["created"])
    updated = datetimestr_to_datetime(info_item["properties"]["updated"])

    status = info_item["properties"]["paii-misc:status"]
    version = info_item["properties"]["paii-misc:version"]

    info = info_item
    geom = json.dumps(info_item["geometry"])

    return {
        "uid": uid,
        "start_datetime": acquired_start,
        "end_datetime": acquired_end,
        "processed": processed,
        "published": published,
        "indexed": indexed,
        "transferred": transferred,
        "archived": archived,
        "created": created,
        "updated": updated,
        "valid_until": None,
        "status": status,
        "version": version,
        "info": info,
        "geom": geom,
    }
